{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydantic import BaseModel\n",
    "from enum import Enum\n",
    "from typing import List, Optional, Union\n",
    "import json\n",
    "import os\n",
    "from typing import Optional\n",
    "from pprint import pprint\n",
    "from autogen_core.models import UserMessage, SystemMessage, AssistantMessage\n",
    "from autogen_ext.models.azure import AzureAIChatCompletionClient\n",
    "from azure.core.credentials import AzureKeyCredential\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "load_dotenv()\n",
    "\n",
    "class AgentEnum(str, Enum):\n",
    "    FlightBooking = \"flight_booking\"\n",
    "    HotelBooking = \"hotel_booking\"\n",
    "    CarRental = \"car_rental\"\n",
    "    ActivitiesBooking = \"activities_booking\"\n",
    "    DestinationInfo = \"destination_info\"\n",
    "    DefaultAgent = \"default_agent\"\n",
    "    GroupChatManager = \"group_chat_manager\"\n",
    "\n",
    "# Travel SubTask Model\n",
    "\n",
    "\n",
    "class TravelSubTask(BaseModel):\n",
    "    task_details: str\n",
    "    assigned_agent: AgentEnum  # we want to assign the task to the agent\n",
    "\n",
    "\n",
    "class TravelPlan(BaseModel):\n",
    "    main_task: str\n",
    "    subtasks: List[TravelSubTask]\n",
    "    is_greeting: bool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "client = AzureAIChatCompletionClient(\n",
    "    model=\"gpt-4o-mini\",\n",
    "    endpoint=\"https://models.inference.ai.azure.com\",\n",
    "    # To authenticate with the model you will need to generate a personal access token (PAT) in your GitHub settings.\n",
    "    # Create your PAT token by following instructions here: https://docs.github.com/en/authentication/keeping-your-account-and-data-secure/managing-your-personal-access-tokens\n",
    "    credential=AzureKeyCredential(os.getenv(\"GITHUB_TOKEN\")),\n",
    "    model_info={\n",
    "        \"json_output\": False,\n",
    "        \"function_calling\": True,\n",
    "        \"vision\": True,\n",
    "        \"family\": \"unknown\",\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'main_task': 'Plan a family trip from Singapore to Melbourne.',\n",
      " 'subtasks': [{'assigned_agent': 'flight_booking',\n",
      "               'task_details': 'Book round-trip flights from Singapore to '\n",
      "                               'Melbourne.'},\n",
      "              {'assigned_agent': 'hotel_booking',\n",
      "               'task_details': 'Book family-friendly accommodations in '\n",
      "                               'Melbourne.'},\n",
      "              {'assigned_agent': 'car_rental',\n",
      "               'task_details': 'Arrange car rental for the family to explore '\n",
      "                               'Melbourne.'},\n",
      "              {'assigned_agent': 'activities_booking',\n",
      "               'task_details': 'Book family-friendly activities and '\n",
      "                               'attractions in Melbourne.'},\n",
      "              {'assigned_agent': 'destination_info',\n",
      "               'task_details': 'Provide information about family-friendly '\n",
      "                               'destinations in Melbourne.'}]}\n"
     ]
    }
   ],
   "source": [
    "# Define the user message\n",
    "messages = [\n",
    "    SystemMessage(content=\"\"\"You are an planner agent.\n",
    "    Your job is to decide which agents to run based on the user's request.\n",
    "                      Provide your response in JSON format with the following structure:\n",
    "{'main_task': 'Plan a family trip from Singapore to Melbourne.',\n",
    " 'subtasks': [{'assigned_agent': 'flight_booking',\n",
    "               'task_details': 'Book round-trip flights from Singapore to '\n",
    "                               'Melbourne.'}\n",
    "    Below are the available agents specialised in different tasks:\n",
    "    - FlightBooking: For booking flights and providing flight information\n",
    "    - HotelBooking: For booking hotels and providing hotel information\n",
    "    - CarRental: For booking cars and providing car rental information\n",
    "    - ActivitiesBooking: For booking activities and providing activity information\n",
    "    - DestinationInfo: For providing information about destinations\n",
    "    - DefaultAgent: For handling general requests\"\"\", source=\"system\"),\n",
    "    UserMessage(\n",
    "        content=\"Create a travel plan for a family of 2 kids from Singapore to Melboune\", source=\"user\"),\n",
    "]\n",
    "\n",
    "response = await client.create(messages=messages, extra_create_args={\"response_format\": 'json_object'})\n",
    "\n",
    "\n",
    "response_content: Optional[str] = response.content if isinstance(\n",
    "    response.content, str) else None\n",
    "if response_content is None:\n",
    "    raise ValueError(\"Response content is not a valid JSON string\" )\n",
    "\n",
    "pprint(json.loads(response_content))\n",
    "\n",
    "# # Ensure the response content is a valid JSON string before loading it\n",
    "# response_content: Optional[str] = response.content if isinstance(\n",
    "#     response.content, str) else None\n",
    "# if response_content is None:\n",
    "#     raise ValueError(\"Response content is not a valid JSON string\")\n",
    "\n",
    "# # Print the response content after loading it as JSON\n",
    "# pprint(json.loads(response_content))\n",
    "\n",
    "# Validate the response content with the MathReasoning model\n",
    "# TravelPlan.model_validate(json.loads(response_content))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
